import os

import numpy as np
from PIL import Image
from scipy.spatial.distance import cosine
from dataset_gen import gen_dataset

images_ori, images_shade = gen_dataset()


def euclidean_distance_matrix(matrix, vector):
    distances = np.linalg.norm(matrix - vector, axis=1)
    min_distance_index = np.argmin(distances)
    return min_distance_index


def search_patch(img, pos, margin=3):
    x, y = pos
    x_min = x - margin if x - margin >= 0 else 0
    x_max = x + margin
    y_min = y - margin if y - margin >= 0 else 0
    y_max = y + margin

    condition = (images_shade[:, 0] >= x_min) & (images_shade[:, 0] <= x_max) \
                & (images_shade[:, 1] >= y_min) & (images_shade[:, 1] <= y_max)
    images_shade_region = images_shade[condition]

    if images_shade_region.shape[0] == 0:
        return img

    images_ori_region = images_ori[condition]

    match = euclidean_distance_matrix(images_shade_region[:, 2:], img)
    return images_ori_region[match][2:]


def sub_images_match(img_path, sub_image_size, step_size):
    print(F'Start restore {img_path}')
    img = Image.open(img_path)
    width, height = img.size

    # 子图分割
    for y in range(0, height, step_size):
        for x in range(0, width, step_size):
            # 计算子图的区域
            region = (x, y, x + sub_image_size, y + sub_image_size)

            if region[2] > width or region[3] > height:
                continue

            # 替换子图
            sub_img = np.array(img.crop(region)).flatten()
            sub_img_restore = search_patch(sub_img, (x, y))
            sub_img_restore = sub_img_restore.reshape((sub_image_size, sub_image_size))
            img.paste(Image.fromarray(sub_img_restore), region)
    return img


def image_restore():
    # test_path = r'./data/202303/dataset/test'
    test_path = r'./data/20230827/test'
    sub_image_size = 5
    step_size = 5

    for file in os.listdir(test_path):
        img_path = os.path.join(test_path, file)
        res = sub_images_match(img_path, sub_image_size, step_size)
        res.save(os.path.join(test_path, F'res_{file}'))
        print(file)


if __name__ == '__main__':
    image_restore()
